#!/usr/bin/env python3
"""
TCD-VoIP 音频质量评估结果分析脚本

此脚本分析音频质量评估结果，包括：
1. 整体RMSE对比
2. 分劣化类型的详细分析
3. 相关性分析
4. 可视化图表生成

作者: 基于课程设计任务书要求  
日期: 2025年9月7日
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import mean_squared_error
from scipy.stats import pearsonr, spearmanr
import warnings
import os
warnings.filterwarnings('ignore')

# 配置中文字体支持
def setup_chinese_fonts():
    """配置matplotlib中文字体支持"""
    global USE_CHINESE_LABELS
    
    try:
        import matplotlib.font_manager as fm
        
        # 直接注册Noto CJK字体
        noto_font_path = '/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc'
        
        if os.path.exists(noto_font_path):
            # 注册字体
            fm.fontManager.addfont(noto_font_path)
            
            # 获取字体属性
            font_prop = fm.FontProperties(fname=noto_font_path)
            font_name = font_prop.get_name()
            
            # 设置为默认字体
            plt.rcParams['font.sans-serif'] = [font_name, 'DejaVu Sans']
            plt.rcParams['axes.unicode_minus'] = False
            
            USE_CHINESE_LABELS = True
            print(f"✓ 成功配置中文字体: {font_name}")
            return True
        else:
            # 备用方案：使用已有字体配置
            plt.rcParams['font.sans-serif'] = ['DejaVu Sans', 'Liberation Sans']
            plt.rcParams['axes.unicode_minus'] = False
            USE_CHINESE_LABELS = False
            print("⚠ Noto字体未找到，使用英文标签")
            return False
            
    except Exception as e:
        USE_CHINESE_LABELS = False
        plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False
        print(f"字体配置错误: {e}，使用英文标签")
        return False

# 定义标签字典（中英文对照）
LABELS = {
    'chinese': {
        'title': '音频质量评估结果分析',
        'rmse': 'RMSE',
        'correlation': '相关系数',
        'degradation_type': '劣化类型',
        'mos_score': 'MOS分数',
        'predicted_score': '预测分数',
        'pesq': 'PESQ',
        'visqol': 'ViSQOL',
        'p563': 'P.563',
        'clip': '裁剪',
        'noise': '噪声',
        'echo': '回声',
        'chop': '切断',
        'competing_speaker': '竞争说话人'
    },
    'english': {
        'title': 'Audio Quality Assessment Results Analysis',
        'rmse': 'RMSE',
        'correlation': 'Correlation',
        'degradation_type': 'Degradation Type',
        'mos_score': 'MOS Score',
        'predicted_score': 'Predicted Score',
        'pesq': 'PESQ',
        'visqol': 'ViSQOL',
        'p563': 'P.563',
        'clip': 'Clip',
        'noise': 'Noise',
        'echo': 'Echo',
        'chop': 'Chop',
        'competing_speaker': 'Competing Speaker'
    }
}

def get_label(key):
    """获取标签（根据字体支持情况选择中英文）"""
    if USE_CHINESE_LABELS and key in LABELS['chinese']:
        return LABELS['chinese'][key]
    return LABELS['english'].get(key, key)

# 全局变量
USE_CHINESE_LABELS = True

# 设置中文字体
setup_chinese_fonts()

# ==================== 配置区 ====================
RESULTS_FILE = 'evaluation_results.csv'
OUTPUT_DIR = 'analysis_results'
FIGURES_DIR = os.path.join(OUTPUT_DIR, 'figures')
TABLES_DIR = os.path.join(OUTPUT_DIR, 'tables')

# 工具名称映射
TOOL_NAMES = {
    'pesq_mos_lqo': 'PESQ',
    'visqol_mos_lqo': 'ViSQOL', 
    'p563_mos_lqo': 'P.563'
}

# 劣化类型颜色映射
DEGRADATION_COLORS = {
    'Chop': '#FF6B6B',
    'Clip': '#4ECDC4', 
    'Competing Speaker': '#45B7D1',
    'Echo': '#96CEB4',
    'Noise': '#FFEAA7'
}

# ==================== 工具函数 ====================
def ensure_directories():
    """确保输出目录存在"""
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    os.makedirs(FIGURES_DIR, exist_ok=True)
    os.makedirs(TABLES_DIR, exist_ok=True)

def load_and_clean_data():
    """加载并清理数据"""
    try:
        df = pd.read_csv(RESULTS_FILE, encoding='utf-8-sig')
        print(f"加载数据: {len(df)} 行")
        
        # 显示数据基本信息
        print(f"\n数据集信息:")
        print(f"- 劣化类型: {df['degradation_type'].unique()}")
        print(f"- MOS-LQS范围: {df['mos_lqs'].min():.2f} - {df['mos_lqs'].max():.2f}")
        
        # 检查每个工具的可用数据
        for tool_col, tool_name in TOOL_NAMES.items():
            available = df[tool_col].notna().sum()
            total = len(df)
            print(f"- {tool_name}: {available}/{total} ({available/total*100:.1f}%)")
        
        return df
    except FileNotFoundError:
        print(f"错误: 找不到结果文件 {RESULTS_FILE}")
        print("请先运行 'python run_evaluation.py'")
        return None

def calculate_metrics(y_true, y_pred):
    """计算各种评估指标"""
    # 移除NaN值
    mask = ~(np.isnan(y_true) | np.isnan(y_pred))
    y_true_clean = y_true[mask]
    y_pred_clean = y_pred[mask]
    
    if len(y_true_clean) < 2:
        return {
            'rmse': np.nan,
            'mae': np.nan,
            'pearson_r': np.nan,
            'spearman_r': np.nan,
            'count': len(y_true_clean)
        }
    
    # RMSE
    rmse = np.sqrt(mean_squared_error(y_true_clean, y_pred_clean))
    
    # MAE  
    mae = np.mean(np.abs(y_true_clean - y_pred_clean))
    
    # Pearson相关系数
    pearson_r, _ = pearsonr(y_true_clean, y_pred_clean)
    
    # Spearman相关系数
    spearman_r, _ = spearmanr(y_true_clean, y_pred_clean)
    
    return {
        'rmse': rmse,
        'mae': mae, 
        'pearson_r': pearson_r,
        'spearman_r': spearman_r,
        'count': len(y_true_clean)
    }

def analyze_overall_performance(df):
    """分析整体性能"""
    print("\n" + "="*60)
    print("整体性能分析")
    print("="*60)
    
    overall_results = []
    
    for tool_col, tool_name in TOOL_NAMES.items():
        if tool_col in df.columns:
            y_true = df['mos_lqs'].values
            y_pred = df[tool_col].values
            
            metrics = calculate_metrics(y_true, y_pred)
            
            overall_results.append({
                'Tool': tool_name,
                'RMSE': metrics['rmse'],
                'MAE': metrics['mae'],
                'Pearson R': metrics['pearson_r'],
                'Spearman R': metrics['spearman_r'],
                'Sample Count': metrics['count']
            })
            
            print(f"\n{tool_name}:")
            print(f"  RMSE: {metrics['rmse']:.4f}")
            print(f"  MAE:  {metrics['mae']:.4f}")
            print(f"  Pearson R: {metrics['pearson_r']:.4f}")
            print(f"  Spearman R: {metrics['spearman_r']:.4f}")
            print(f"  样本数: {metrics['count']}")
    
    # 保存整体结果表格
    overall_df = pd.DataFrame(overall_results)
    overall_df.to_csv(os.path.join(TABLES_DIR, 'overall_performance.csv'), 
                     index=False, encoding='utf-8-sig')
    
    return overall_df

def analyze_by_degradation(df):
    """按劣化类型分析性能"""
    print("\n" + "="*60)
    print("分劣化类型性能分析")
    print("="*60)
    
    degradation_results = []
    
    for deg_type in df['degradation_type'].unique():
        if pd.isna(deg_type):
            continue
            
        subset = df[df['degradation_type'] == deg_type]
        print(f"\n{deg_type} ({len(subset)} 样本):")
        
        for tool_col, tool_name in TOOL_NAMES.items():
            if tool_col in subset.columns:
                y_true = subset['mos_lqs'].values
                y_pred = subset[tool_col].values
                
                metrics = calculate_metrics(y_true, y_pred)
                
                degradation_results.append({
                    'Degradation Type': deg_type,
                    'Tool': tool_name,
                    'RMSE': metrics['rmse'],
                    'MAE': metrics['mae'],
                    'Pearson R': metrics['pearson_r'],
                    'Spearman R': metrics['spearman_r'],
                    'Sample Count': metrics['count']
                })
                
                print(f"  {tool_name}: RMSE={metrics['rmse']:.4f}, "
                      f"MAE={metrics['mae']:.4f}, R={metrics['pearson_r']:.4f}")
    
    # 保存分类结果表格
    degradation_df = pd.DataFrame(degradation_results)
    degradation_df.to_csv(os.path.join(TABLES_DIR, 'degradation_performance.csv'),
                         index=False, encoding='utf-8-sig')
    
    return degradation_df

def create_overall_comparison_plot(overall_df):
    """创建整体性能对比图"""
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle('整体性能对比', fontsize=16, fontweight='bold')
    
    # RMSE对比
    axes[0,0].bar(overall_df['Tool'], overall_df['RMSE'], 
                  color=['#FF6B6B', '#4ECDC4', '#45B7D1'])
    axes[0,0].set_title('RMSE 对比 (越低越好)')
    axes[0,0].set_ylabel('RMSE')
    axes[0,0].grid(axis='y', alpha=0.3)
    
    # MAE对比  
    axes[0,1].bar(overall_df['Tool'], overall_df['MAE'],
                  color=['#FF6B6B', '#4ECDC4', '#45B7D1'])
    axes[0,1].set_title('MAE 对比 (越低越好)')
    axes[0,1].set_ylabel('MAE')
    axes[0,1].grid(axis='y', alpha=0.3)
    
    # Pearson相关系数
    axes[1,0].bar(overall_df['Tool'], overall_df['Pearson R'],
                  color=['#FF6B6B', '#4ECDC4', '#45B7D1'])
    axes[1,0].set_title('Pearson 相关系数 (越高越好)')
    axes[1,0].set_ylabel('Pearson R')
    axes[1,0].set_ylim(0, 1)
    axes[1,0].grid(axis='y', alpha=0.3)
    
    # Spearman相关系数
    axes[1,1].bar(overall_df['Tool'], overall_df['Spearman R'],
                  color=['#FF6B6B', '#4ECDC4', '#45B7D1'])
    axes[1,1].set_title('Spearman 相关系数 (越高越好)')
    axes[1,1].set_ylabel('Spearman R')
    axes[1,1].set_ylim(0, 1)
    axes[1,1].grid(axis='y', alpha=0.3)
    
    # 为每个子图添加数值标签
    for ax in axes.flat:
        for i, bar in enumerate(ax.patches):
            height = bar.get_height()
            if not np.isnan(height):
                ax.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                       f'{height:.3f}', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.savefig(os.path.join(FIGURES_DIR, 'overall_performance_comparison.png'), 
                dpi=300, bbox_inches='tight')
    plt.show()

def create_degradation_comparison_plot(degradation_df):
    """创建分劣化类型的性能对比图"""
    # RMSE分劣化类型对比
    pivot_rmse = degradation_df.pivot(index='Degradation Type', 
                                      columns='Tool', values='RMSE')
    
    fig, axes = plt.subplots(2, 2, figsize=(16, 12))
    fig.suptitle('分劣化类型性能分析', fontsize=16, fontweight='bold')
    
    # RMSE热力图
    sns.heatmap(pivot_rmse, annot=True, fmt='.3f', cmap='YlOrRd',
                ax=axes[0,0], cbar_kws={'label': 'RMSE'})
    axes[0,0].set_title('RMSE 热力图')
    axes[0,0].set_xlabel('评估工具')
    axes[0,0].set_ylabel('劣化类型')
    
    # RMSE柱状图
    pivot_rmse.plot(kind='bar', ax=axes[0,1], width=0.8)
    axes[0,1].set_title('RMSE 柱状对比')
    axes[0,1].set_ylabel('RMSE')
    axes[0,1].legend(title='工具')
    axes[0,1].tick_params(axis='x', rotation=45)
    axes[0,1].grid(axis='y', alpha=0.3)
    
    # Pearson相关系数
    pivot_pearson = degradation_df.pivot(index='Degradation Type', 
                                        columns='Tool', values='Pearson R')
    sns.heatmap(pivot_pearson, annot=True, fmt='.3f', cmap='RdYlBu_r',
                ax=axes[1,0], cbar_kws={'label': 'Pearson R'})
    axes[1,0].set_title('Pearson 相关系数热力图')
    axes[1,0].set_xlabel('评估工具')
    axes[1,0].set_ylabel('劣化类型')
    
    # 样本数分布
    pivot_count = degradation_df.pivot(index='Degradation Type', 
                                      columns='Tool', values='Sample Count')
    pivot_count.plot(kind='bar', ax=axes[1,1], width=0.8)
    axes[1,1].set_title('样本数分布')
    axes[1,1].set_ylabel('样本数')
    axes[1,1].legend(title='工具')
    axes[1,1].tick_params(axis='x', rotation=45)
    axes[1,1].grid(axis='y', alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(os.path.join(FIGURES_DIR, 'degradation_type_analysis.png'), 
                dpi=300, bbox_inches='tight')
    plt.show()

def create_scatter_plots(df):
    """创建散点图显示主观vs客观分数的关系"""
    available_tools = [col for col in TOOL_NAMES.keys() if col in df.columns and df[col].notna().sum() > 0]
    n_tools = len(available_tools)
    
    if n_tools == 0:
        print("警告: 没有可用的工具数据用于散点图")
        return
    
    fig, axes = plt.subplots(1, n_tools, figsize=(6*n_tools, 5), sharey=True)
    if n_tools == 1:
        axes = [axes]
        
    fig.suptitle('主观 MOS vs 客观 MOS 散点图', fontsize=16, fontweight='bold')
    
    for i, tool_col in enumerate(available_tools):
        tool_name = TOOL_NAMES[tool_col]
        
        # 过滤掉NaN值
        mask = df[tool_col].notna() & df['mos_lqs'].notna()
        if mask.sum() == 0:
            continue
            
        subset = df[mask]
        
        # 按劣化类型着色
        for deg_type in subset['degradation_type'].unique():
            if pd.isna(deg_type):
                continue
            deg_subset = subset[subset['degradation_type'] == deg_type]
            color = DEGRADATION_COLORS.get(deg_type, '#333333')
            
            axes[i].scatter(deg_subset['mos_lqs'], deg_subset[tool_col], 
                          c=color, label=deg_type, alpha=0.7, s=30)
        
        # 添加理想线 y=x
        min_val = min(subset['mos_lqs'].min(), subset[tool_col].min())
        max_val = max(subset['mos_lqs'].max(), subset[tool_col].max())
        axes[i].plot([min_val, max_val], [min_val, max_val], 
                    'r--', alpha=0.8, linewidth=2, label='理想线 (y=x)')
        
        # 计算并显示相关系数
        pearson_r, _ = pearsonr(subset['mos_lqs'], subset[tool_col])
        rmse = np.sqrt(mean_squared_error(subset['mos_lqs'], subset[tool_col]))
        
        axes[i].set_title(f'{tool_name}\nR={pearson_r:.3f}, RMSE={rmse:.3f}')
        axes[i].set_xlabel('主观 MOS (MOS-LQS)')
        if i == 0:
            axes[i].set_ylabel('客观 MOS (MOS-LQO)')
        axes[i].grid(True, alpha=0.3)
        axes[i].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        
        # 设置轴范围
        axes[i].set_xlim(0.5, 5.5)
        axes[i].set_ylim(0.5, 5.5)
    
    plt.tight_layout()
    plt.savefig(os.path.join(FIGURES_DIR, 'subjective_vs_objective_scatter.png'), 
                dpi=300, bbox_inches='tight')
    plt.show()

def create_distribution_plots(df):
    """创建分数分布图"""
    available_tools = [col for col in TOOL_NAMES.keys() if col in df.columns and df[col].notna().sum() > 0]
    
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    fig.suptitle('MOS 分数分布分析', fontsize=16, fontweight='bold')
    
    # 主观MOS分布
    axes[0,0].hist(df['mos_lqs'].dropna(), bins=20, alpha=0.7, color='skyblue', edgecolor='black')
    axes[0,0].set_title('主观 MOS (MOS-LQS) 分布')
    axes[0,0].set_xlabel('MOS-LQS')
    axes[0,0].set_ylabel('频次')
    axes[0,0].grid(axis='y', alpha=0.3)
    
    # 客观MOS分布对比
    for tool_col in available_tools:
        tool_name = TOOL_NAMES[tool_col]
        data = df[tool_col].dropna()
        if len(data) > 0:
            axes[0,1].hist(data, bins=20, alpha=0.5, label=tool_name, edgecolor='black')
    
    axes[0,1].set_title('客观 MOS (MOS-LQO) 分布对比')
    axes[0,1].set_xlabel('MOS-LQO')
    axes[0,1].set_ylabel('频次')
    axes[0,1].legend()
    axes[0,1].grid(axis='y', alpha=0.3)
    
    # 按劣化类型的MOS分布
    deg_types = df['degradation_type'].unique()
    deg_types = [dt for dt in deg_types if not pd.isna(dt)]
    
    for deg_type in deg_types[:5]:  # 最多显示5种劣化类型
        subset = df[df['degradation_type'] == deg_type]
        color = DEGRADATION_COLORS.get(deg_type, '#333333')
        axes[1,0].hist(subset['mos_lqs'].dropna(), bins=15, alpha=0.6, 
                      label=deg_type, color=color, edgecolor='black')
    
    axes[1,0].set_title('按劣化类型的主观 MOS 分布')
    axes[1,0].set_xlabel('MOS-LQS')
    axes[1,0].set_ylabel('频次')
    axes[1,0].legend()
    axes[1,0].grid(axis='y', alpha=0.3)
    
    # 误差分布
    if available_tools:
        for tool_col in available_tools:
            tool_name = TOOL_NAMES[tool_col]
            mask = df[tool_col].notna() & df['mos_lqs'].notna()
            if mask.sum() > 0:
                errors = df.loc[mask, tool_col] - df.loc[mask, 'mos_lqs']
                axes[1,1].hist(errors, bins=20, alpha=0.5, label=f'{tool_name} 误差', edgecolor='black')
    
    axes[1,1].set_title('预测误差分布')
    axes[1,1].set_xlabel('预测误差 (MOS-LQO - MOS-LQS)')
    axes[1,1].set_ylabel('频次')
    axes[1,1].legend()
    axes[1,1].grid(axis='y', alpha=0.3)
    axes[1,1].axvline(x=0, color='red', linestyle='--', alpha=0.8)
    
    plt.tight_layout()
    plt.savefig(os.path.join(FIGURES_DIR, 'score_distributions.png'), 
                dpi=300, bbox_inches='tight')
    plt.show()

def generate_summary_report(overall_df, degradation_df):
    """生成总结报告"""
    report = []
    report.append("# TCD-VoIP 音频质量评估分析报告")
    report.append("")
    
    # 检查是否有有效数据
    if overall_df.empty or overall_df['RMSE'].isna().all():
        report.append("## 警告: 没有有效的评估数据")
        report.append("")
        report.append("可能的原因:")
        report.append("- 音频质量评估工具未正常运行")
        report.append("- 音频文件格式不兼容")
        report.append("- 文件路径错误")
        report.append("")
        report.append("建议:")
        report.append("1. 检查音频文件是否为16kHz WAV格式")
        report.append("2. 验证工具是否能独立运行")
        report.append("3. 检查文件路径和权限")
    else:
        report.append("## 1. 整体性能排名")
        report.append("")
        
        # 按RMSE排序
        valid_data = overall_df.dropna(subset=['RMSE'])
        if not valid_data.empty:
            overall_sorted = valid_data.sort_values('RMSE')
            report.append("### 按RMSE排名 (越低越好):")
            for i, (_, row) in enumerate(overall_sorted.iterrows(), 1):
                report.append(f"{i}. **{row['Tool']}**: RMSE = {row['RMSE']:.4f}")
        
        report.append("")
        
        # 按相关系数排序  
        valid_data_r = overall_df.dropna(subset=['Pearson R'])
        if not valid_data_r.empty:
            overall_sorted_r = valid_data_r.sort_values('Pearson R', ascending=False)
            report.append("### 按Pearson相关系数排名 (越高越好):")
            for i, (_, row) in enumerate(overall_sorted_r.iterrows(), 1):
                report.append(f"{i}. **{row['Tool']}**: R = {row['Pearson R']:.4f}")
        
        report.append("")
        report.append("## 2. 最佳工具推荐")
        report.append("")
        
        # 找出各劣化类型下表现最好的工具
        deg_types = degradation_df['Degradation Type'].dropna().unique()
        for deg_type in deg_types:
            subset = degradation_df[
                (degradation_df['Degradation Type'] == deg_type) & 
                (degradation_df['RMSE'].notna())
            ]
            if not subset.empty:
                best_idx = subset['RMSE'].idxmin()
                best_tool = subset.loc[best_idx, 'Tool']
                best_rmse = subset.loc[best_idx, 'RMSE']
                report.append(f"- **{deg_type}** 劣化: {best_tool} (RMSE = {best_rmse:.4f})")
        
        report.append("")
        report.append("## 3. 关键发现")
        report.append("")
        
        # 一些自动化的发现
        if not valid_data.empty:
            best_overall = valid_data.loc[valid_data['RMSE'].idxmin(), 'Tool']
            worst_overall = valid_data.loc[valid_data['RMSE'].idxmax(), 'Tool']
            
            report.append(f"- 整体表现最佳工具: **{best_overall}**")
            report.append(f"- 整体表现最差工具: **{worst_overall}**")
    
    # 保存报告
    with open(os.path.join(OUTPUT_DIR, 'analysis_summary.md'), 'w', encoding='utf-8') as f:
        f.write('\n'.join(report))
    
    print("\n" + "="*60)
    print("分析总结")
    print("="*60)
    for line in report:
        if not line.startswith('#'):
            print(line)

# ==================== 主函数 ====================
def main():
    """主函数"""
    print("=" * 60)
    print("TCD-VoIP 音频质量评估结果分析系统")
    print("=" * 60)
    
    # 确保输出目录存在
    ensure_directories()
    
    # 加载数据
    print("\n1. 加载评估结果...")
    df = load_and_clean_data()
    if df is None:
        return
    
    # 整体性能分析
    print("\n2. 整体性能分析...")
    overall_df = analyze_overall_performance(df)
    
    # 分劣化类型分析
    print("\n3. 分劣化类型分析...")
    degradation_df = analyze_by_degradation(df)
    
    # 生成图表
    print("\n4. 生成可视化图表...")
    
    print("  - 整体性能对比图...")
    create_overall_comparison_plot(overall_df)
    
    print("  - 分劣化类型分析图...")
    create_degradation_comparison_plot(degradation_df)
    
    print("  - 散点图分析...")
    create_scatter_plots(df)
    
    print("  - 分布图分析...")
    create_distribution_plots(df)
    
    # 生成总结报告
    print("\n5. 生成分析报告...")
    generate_summary_report(overall_df, degradation_df)
    
    print(f"\n分析完成！")
    print(f"结果保存在: {OUTPUT_DIR}/")
    print(f"- 表格: {TABLES_DIR}/")
    print(f"- 图片: {FIGURES_DIR}/")
    print(f"- 报告: {OUTPUT_DIR}/analysis_summary.md")

if __name__ == "__main__":
    main()
